package cn.iocoder.yudao.framework.mybatis.interceptor;

import cn.hutool.core.util.RandomUtil;
import cn.iocoder.yudao.framework.datasource.config.DynamicRoutingDataSource;
import cn.iocoder.yudao.framework.datasource.local.ThreadLocalContext;
import cn.iocoder.yudao.framework.msdb.ICustomizeRoutingDataSource;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * TODO
 * 读写分离拦截器
 * @className: MsdbInterceptor
 * @author: tangkh
 * @date: 2023-09-19 17:50
 */
@Intercepts({@Signature(
        type = Executor.class,
        method = "update",
        args = {MappedStatement.class, Object.class}
), @Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}
), @Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
)})
public class MsdbInterceptor implements Interceptor {
    private static final Logger log = LoggerFactory.getLogger(MsdbInterceptor.class);
    private static final String READ_ONLY_REGEX = "[\\s\\S]*.*#query=ReAdOnLY#.*[\\s\\S]*";
    @Value("${spring.datasource.dynamic.primary}")
    private String primary;
    @Autowired(
            required = false
    )
    private ICustomizeRoutingDataSource contextPathInterceptor;
    @Value("${spring.datasource.dynamic.extends.multi-tenant.enable:false}")
    private boolean multiTenantEnable;
    @Value("${spring.datasource.dynamic.extends.dbkey:}")
    private String defaultDbkey;

    public MsdbInterceptor() {
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        return invocation.proceed();
    }

    private String getMaster() {
        String master = DynamicDataSourceContextHolder.peek();
        String dbKey = (String) ThreadLocalContext.getAttachment("dbKey");
        if (StringUtils.hasText(dbKey)) {
            master = dbKey;
        } else if (this.contextPathInterceptor != null) {
            String value = this.contextPathInterceptor.getDbKey();
            if (StringUtils.hasText(value)) {
                master = value;
            } else {
                master = StringUtils.hasText(master) ? master : this.defaultDbkey;
            }
        }

        if (this.multiTenantEnable) {
            RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
            if (requestAttributes != null) {
                HttpServletRequest request = ((ServletRequestAttributes)requestAttributes).getRequest();
                String orgId = request.getParameter("gwOrgId");
                master = this.getDatasourceName(master, orgId);
            } else {
                String orgId = (String)ThreadLocalContext.getAttachment("orgId");
                master = this.getDatasourceName(master, orgId);
            }
        }

        return master;
    }

    private void chooseDB(String master, String slave, Object[] objects, MappedStatement ms) {
        if (ms.getSqlCommandType().equals(SqlCommandType.SELECT)) {
            BoundSql boundSql = ms.getSqlSource().getBoundSql(objects[1]);
            String sql = boundSql.getSql();
            if (sql.matches("[\\s\\S]*.*#query=ReAdOnLY#.*[\\s\\S]*")) {
                this.chooseReadDb(master, slave);
            } else {
                log.info("走主库,数据库名:{}", master);
                DynamicDataSourceContextHolder.push(master);
            }
        } else {
            log.info("走主库,数据库名:{}", master);
            DynamicDataSourceContextHolder.push(master);
        }

    }

    private void chooseReadDb(String master, String slave) {
        List<String> readDbList = new ArrayList();
        DynamicRoutingDataSource dynamicRoutingDataSource = null;// todo (DynamicRoutingDataSource)SpringContextUtils.getBean(DynamicRoutingDataSource.class);
        Map<String, DataSource> dataSources = dynamicRoutingDataSource.getDataSources();
        DataSource slaveDs = (DataSource)dataSources.get(slave);
        if (slaveDs != null) {
            readDbList.add(slave);
        }

        if (StringUtils.hasLength(master) && master.contains("_slave")) {
            String[] slaves = master.split("_slave");
            readDbList.add(slaves[0]);
        } else {
            readDbList.add(master);
        }

        int randomInt = RandomUtil.randomInt(10, 100);
        int index = randomInt % readDbList.size();
        String readDbsId = (String)readDbList.get(index);
        DynamicDataSourceContextHolder.push(readDbsId);
        log.info("读写分离，走{}库，数据库名: {}, 序号: {}", new Object[]{readDbsId.contains("_slave") ? "读" : "写", readDbsId, index});
    }

    private String getDatasourceName(String key, String orgId) {
        if (!Objects.isNull(orgId) && !key.contains(orgId)) {
            String substring = "".equals(key) ? this.primary : key;
            return orgId + "_" + substring;
        } else {
            return StringUtils.isEmpty(key) ? this.primary : key;
        }
    }
}
